package me.zhengjie.controller;

import io.prestosql.client.ClientException;
import me.zhengjie.common.DlaServerConfig;
import me.zhengjie.exception.QueryErrorException;
import me.zhengjie.result.PrestoQueryResult;
import me.zhengjie.service.PrestoService;
import me.zhengjie.util.MetadataUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RestController;

import javax.annotation.Resource;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.stream.Collectors;

import static java.lang.String.format;
import static me.zhengjie.util.AccessControlUtil.sendForbiddenError;
import static me.zhengjie.util.Constants.YANAGISHIMA_COMMENT;
import static me.zhengjie.util.HttpRequestUtil.getRequiredParameter;
import static me.zhengjie.util.JsonUtil.writeJSON;


@RestController
public class PrestoAuthorController {
    private static final Logger LOGGER = LoggerFactory.getLogger(PrestoAuthorController.class);

    @Resource
    private PrestoService prestoService;
    @Resource
    private DlaServerConfig config;

    @PostMapping("/presto")
    public Object presto(HttpServletRequest request, HttpServletResponse response){
        Map<String, Object> responseBody = new HashMap<>();
        try {
            Optional<String> queryOptional = Optional.ofNullable(request.getParameter("query"));
            queryOptional.ifPresent(query -> {

                String dataSource =getRequiredParameter(request, "dataSource");
                String userName = getUsername(dataSource);
                Optional<String> prestoUser = Optional.ofNullable(userName);
                Optional<String> prestoPassword = Optional.ofNullable(request.getParameter("password"));
                if (config.isUserRequired() && userName == null) {
                    sendForbiddenError(response);
                    return;
                }
                try {
                    String coordinatorServer = config.getPrestoCoordinatorServerOrNull(dataSource);
                    if (coordinatorServer == null) {
                        writeJSON(response, responseBody);
                        return ;
                    }
                    if (userName != null) {
                        LOGGER.info(format("%s executed %s in %s", userName, query, dataSource));
                    }
                    boolean storeFlag = Boolean.parseBoolean(Optional.ofNullable(request.getParameter("store")).orElse("false"));
                    if (prestoUser.isPresent() && prestoPassword.isPresent()) {
                        if (prestoUser.get().length() == 0) {
                            responseBody.put("error", "user is empty");
                            writeJSON(response, responseBody);
                            return;
                        }
                    }
                    PrestoQueryResult prestoQueryResult;
                    prestoQueryResult = prestoService.doQuery(dataSource, query, userName, prestoUser, prestoPassword, storeFlag, Integer.MAX_VALUE);
                    String queryId = prestoQueryResult.getQueryId();
                    responseBody.put("queryId", queryId);
                    if (prestoQueryResult.getUpdateType() == null) {
                        responseBody.put("headers", prestoQueryResult.getColumns());

                        if (query.startsWith(YANAGISHIMA_COMMENT + "SHOW SCHEMAS FROM")) {
                            String catalog = query.substring((YANAGISHIMA_COMMENT + "SHOW SCHEMAS FROM").length()).trim();
                            List<String> invisibleSchemas = config.getInvisibleSchemas(dataSource, catalog);
                            responseBody.put("results", prestoQueryResult.getRecords().stream().filter(list -> !invisibleSchemas.contains(list.get(0))).collect(Collectors.toList()));
                        } else {
                            responseBody.put("results", prestoQueryResult.getRecords());
                        }
                        responseBody.put("lineNumber", Long.toString(prestoQueryResult.getLineNumber()));
                        responseBody.put("rawDataSize", prestoQueryResult.getRawDataSize().toString());
                        Optional<String> warningMessageOptinal = Optional.ofNullable(prestoQueryResult.getWarningMessage());
                        warningMessageOptinal.ifPresent(warningMessage -> {
                            responseBody.put("warn", warningMessage);
                        });
                        if (query.startsWith(YANAGISHIMA_COMMENT + "DESCRIBE") && config.getMetadataServiceUrl(dataSource).isPresent()) {
                            String[] strings = query.substring(YANAGISHIMA_COMMENT.length() + "DESCRIBE ".length()).split("\\.");
                            String schema = strings[1];
                            String table = strings[2].substring(1, strings[2].length() - 1);
                            MetadataUtil.setMetadata(config.getMetadataServiceUrl(dataSource).get(), responseBody, schema, table, prestoQueryResult.getRecords());
                        }
                    }
                } catch (QueryErrorException e) {
                    LOGGER.warn(e.getCause().getMessage());
                    responseBody.put("error", e.getCause().getMessage());
                    responseBody.put("queryId", e.getQueryId());
                } catch (ClientException e) {
                    if (prestoUser.isPresent()) {
                        LOGGER.error(format("%s failed to be authenticated", prestoUser.get()));
                    }
                    LOGGER.error(e.getMessage(), e);
                    responseBody.put("error", e.getMessage());
                } catch (Throwable e) {
                    LOGGER.error(e.getMessage(), e);
                    responseBody.put("error", e.getMessage());
                }
            });
        } catch (Throwable e) {
            LOGGER.error(e.getMessage(), e);
            responseBody.put("error", e.getMessage());
        }
        return responseBody;
    }

    private String getUsername(String dataSource) {
        String user = config.getUser(dataSource);
        if (user != null ) {
            return user;
        }
        return null;
    }

}
